# -*- coding: utf-8 -*-
"""
Created on Fri Nov 26 11:15:44 2021

@author: bwm
"""
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
class strategy:
        #####此区域写我们的交易策略#################################################################################
    """
    策略：应用随机森林算法选出最有可能上涨的股票。
        每天进行一次随机森林计算，判断股票池内每只股票在当天上涨的概率，卖出所有存量股票并选取上涨概率最高的三支股票买入。
        如果最高的三个上涨概率都小于0.55，则当天不买入任何股票。
        训练集：2020年1月1日开始的，到交易日前两天为止的历史交易信息。
        预测集：交易日前一天的交易数据。通过前一天的数据预测；股票当天是否会上涨。
        训练集选用的特征：历史交易数据的每日开盘价、收盘价、市盈率、市净率、换手率、成交量等共15个特征。
        训练标签：相比前一天上涨为1，下跌为-1。
        
    """
    def __init__(self, date_all, me, buy_list, find,data_all,start_date):
        self.date_all = date_all
        self.me = me
        self.buy_list = buy_list
        self.find = find
        self.data_all = data_all
        self.start_date = start_date
    
    #量价双金叉策略
    def get_strategy_golden_price_num(self,date):
                
        golden_buy_list = []  
        
        goldenprice_buy_list = []    
        for stock in self.buy_list:
            averageprice_last5d = 0
            for i in range(1,6):
                date_i = self.find.get_date(date,-i)
                averageprice_last5d += self.find.get_data(date_i,stock,'收盘价')
            averageprice_last5d =averageprice_last5d/5    
            averageprice_last10d = 0
            for i in range(1,11):
                date_i = self.find.get_date(date,-i)
                averageprice_last10d += self.find.get_data(date_i,stock,'收盘价') 
            averageprice_last10d = averageprice_last10d/10
            averageprice_5d = 0
            for i in range(5):
                date_i = self.find.get_date(date,-i)
                averageprice_5d += self.find.get_data(date_i,stock,'收盘价')
            averageprice_5d =averageprice_5d/5
            averageprice_10d = 0
            for i in range(10):
                date_i = self.find.get_date(date,-i)
                averageprice_10d += self.find.get_data(date_i,stock,'收盘价') 
            averageprice_10d = averageprice_10d/10
            #5日均线上穿10日均线
            if averageprice_5d > averageprice_10d and averageprice_last5d < averageprice_last10d:
                goldenprice_buy_list += [stock]
            
        goldennum_buy_list = []        
        for stock in self.buy_list:
            averagenum_last5d = 0
            for i in range(1,6):
                date_i = self.find.get_date(date,-i)
                averagenum_last5d += self.find.get_data(date_i,stock,'成交量(股)')
            averagenum_last5d =averagenum_last5d/5
            averagenum_last10d = 0
            for i in range(1,11):
                date_i = self.find.get_date(date,-i)
                averagenum_last10d += self.find.get_data(date_i,stock,'成交量(股)') 
            averagenum_last10d = averagenum_last10d/10
            averagenum_5d = 0
            for i in range(5):
                date_i = self.find.get_date(date,-i)
                averagenum_5d += self.find.get_data(date_i,stock,'成交量(股)')
            averagenum_5d =averagenum_5d/5
            averagenum_10d = 0
            for i in range(10):
                date_i = self.find.get_date(date,-i)
                averagenum_10d += self.find.get_data(date_i,stock,'成交量(股)') 
            averagenum_10d = averagenum_10d/10
            #5日均量上穿10日均量
            if averagenum_5d > averagenum_10d and averagenum_last5d < averagenum_last10d:
                goldennum_buy_list += [stock]
        
        golden_buy_list = [stock for stock in goldenprice_buy_list if stock in goldennum_buy_list] 
        
        
        # 每2天轮动换仓
        i = 2
        #金股池没有股票的情况
        if len(golden_buy_list) == 0:
            return
        #金股池只有一只股票的情况
        elif len(golden_buy_list) == 1 and self.date_all.index(date) % i == 0: 
            self.me.sell_all(date)     #空仓卖出
            stock_buy = golden_buy_list[0]   #要买入的股票代码
            price_stock_buy = self.find.get_data(date,stock_buy,'收盘价') #要买入的股票价格
            num = (self.me.cash-1000)/100 // price_stock_buy*100  #要买入的股票数量，减一千是为了留有一定资金避免现金为负，预留手续费
            self.me.buy(stock_buy,num,date)      #买入
        #金股池只有两只股票的情况    
        elif len(golden_buy_list) == 2 and self.date_all.index(date) % i == 0: 
            self.me.sell_all(date)
            stock_buy1 , stock_buy2 = golden_buy_list[0] , golden_buy_list[1]
            price_stock_buy1 , price_stock_buy2 = self.find.get_data(date,stock_buy1,'收盘价') , self.find.get_data(date,stock_buy2,'收盘价')
            num1 , num2 = (self.me.cash-1000)/2/100 // price_stock_buy1*100 , (self.me.cash-1000)/2/100 // price_stock_buy2*100  
            self.me.buy(stock_buy1,num1,date)
            self.me.buy(stock_buy2,num2,date)
        #金股池有三只及以上股票的情况    
        elif self.date_all.index(date) % i == 0:   
            self.me.sell_all(date)
            date_i = self.find.get_date(date,-i)
            ##下面为按上个操作日至今涨幅升序排序金股池股票
            stock_sorting = pd.DataFrame(columns=['涨幅'],index = [])
            for stock in golden_buy_list:
                price_i = self.find.get_data(date_i,stock,'收盘价')
                price_today = self.find.get_data(date,stock,'收盘价')
                rise = price_today/price_i - 1
                stock_sorting.loc[stock] = {'涨幅':rise}
            res = stock_sorting.sort_values(by='涨幅', ascending = True)
            ##
            stock_buy1 , stock_buy2 , stock_buy3 =  res.index[0] ,  res.index[1] ,  res.index[2]
            price_stock_buy1 , price_stock_buy2 , price_stock_buy3 = self.find.get_data(date,stock_buy1,'收盘价') , self.find.get_data(date,stock_buy2,'收盘价') , self.find.get_data(date,stock_buy3,'收盘价')
            num1 , num2 ,num3 = (self.me.cash-1000)/3/100 // price_stock_buy1*100 , (self.me.cash-1000)/3/100 // price_stock_buy2*100 , (self.me.cash-1000)/3/100 // price_stock_buy3*100
            self.me.buy(stock_buy1,num1,date)
            self.me.buy(stock_buy2,num2,date)
            self.me.buy(stock_buy3,num3,date)
            
    #################ETF轮动策略#################################       
    def get_strategy_ETFchange(self,date):
        #6个etf，每个周期排一次序，比如五天，然后买涨幅最大的那个。
        i = 5#周期
        if self.date_all.index(date) % i == 0:
            self.me.sell_all(date)
            date_i = self.find.get_date(date,-i)
            stock_sorting = pd.DataFrame(columns=['涨幅'],index = [])
            for stock in self.buy_list:
                price_i = self.find.get_data(date_i,stock,'收盘价')
                price_today = self.find.get_data(date,stock,'收盘价')
                rise = price_today/price_i - 1
                stock_sorting.loc[stock] = {'涨幅':rise}
            res = stock_sorting.sort_values(by='涨幅', ascending=False)
            stock_buy = res.index[0]
            price_stock_buy = self.find.get_data(date,stock_buy,'收盘价')
            #计算数量这里取了100的倍数
            num = (self.me.cash-1000)/100 // price_stock_buy*100#这里减一千是为了留有一定资金避免现金为负，预留手续费
            self.me.buy(stock_buy,num,date)    
            
    ##################随机森林策略#################################################################################################   
    def get_strategy_Randomforest(self,date):
        # 用上半年作为训练集，下半年作为测试集，选20支龙头股票，预测涨的时候买入，跌的时候卖出。每三天换仓。
        if date in self.date_all[:2]:
            return
        prob_list = []
        base_list = []  # 可选20只股票的历史交易信息
        stock_list = list(self.data_all['2020-01-16'].index)
        for i in range(20):
            base_list.append(self.find.get_hist_data('2020-01-02',date, stock_list[i]))
        for i in range(20):
            base_list[i] = base_list[i].rename(columns = {'开盘价': 'open', '收盘价' : 'close', '时间' : 'date'})
            base_list[i]['close-open'] = (base_list[i]['close'] - base_list[i]['open']) / base_list[i]['open']
            base_list[i]['pre-close'] =  base_list[i]['close'].shift(1)
            base_list[i]['price-change'] = base_list[i]['close'] - base_list[i]['pre-close']
            base_list[i]['p-change'] = base_list[i]['price-change'] / base_list[i]['pre-close'] * 100
            base_list[i] = base_list[i].dropna(how="any")

            
            X = base_list[i][['open','close', '涨跌幅(%)', '成交量(股)','成交笔数', '成交金额','总股本(股)', '总市值', '市盈率(TTM)', 'PE市盈率', 'PB市净率', 'PS市销率', 'PCF市现率', '换手率(%)', 'price-change']]
            y = np.where(base_list[i]['price-change'].shift(-1) > 0, 1, -1) # 标签

            
            X_train, X_test = X[:-1], X[-1:]
            y_train, y_test = y[:-1], y[-1:]

            model = RandomForestClassifier(max_depth=3, n_estimators=10, min_samples_leaf=10, random_state=1)
            model.fit(X_train, y_train)
            y_pred_proba = model.predict_proba(X_test)
            # print(y_pred_proba)
            # print(type(y_pred_proba))
            # print(y_pred_proba[0][1])
            prob_list.append(y_pred_proba[0][1])
            # print(prob_list)
        prob_list.sort(reverse= True)   # 排序
        # print(prob_list)
        buy_index_list = [] # 获取最有可能上涨的三支股票的index
        for i in range(3):
            buy_index_list.append(prob_list.index(prob_list[i]))
        forest_buy_list = [stock_list[i] for i in buy_index_list]    
        if prob_list[2] > 0.55:
            self.me.sell_all(date)     #空仓卖出

            stock_buy1, stock_buy2, stock_buy3 = forest_buy_list[0], forest_buy_list[1], forest_buy_list[2]
            price_stock_buy1 , price_stock_buy2 , price_stock_buy3 = self.find.get_data(date,stock_buy1,'收盘价') , self.find.get_data(date,stock_buy2,'收盘价') , self.find.get_data(date,stock_buy3,'收盘价')
            num1 , num2 ,num3 = (self.me.cash-1000)/3/100 // price_stock_buy1*100 , (self.me.cash-1000)/3/100 // price_stock_buy2*100 , (self.me.cash-1000)/3/100 // price_stock_buy3*100
            self.me.buy(stock_buy1,num1,date)
            self.me.buy(stock_buy2,num2,date)
            self.me.buy(stock_buy3,num3,date)

            # stock_buy = forest_buy_list[0]   #要买入的股票代码
            # price_stock_buy = find.get_data(date,stock_buy,'收盘价') #要买入的股票价格
            # num = (me.cash-1000)/100 // price_stock_buy*100  #要买入的股票数量，减一千是为了留有一定资金避免现金为负，预留手续费
            # me.buy(stock_buy,num)
           
        else:
            self.me.sell_all(date)
           
                   
    ###################################################################################################################